package com.example.backend.service;

import com.example.backend.config.LlmConfig;
import com.example.backend.dto.LlmDto;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.http.HttpHeaders;
import org.springframework.http.MediaType;
import org.springframework.stereotype.Service;
import org.springframework.web.reactive.function.client.WebClient;
import reactor.core.publisher.Mono;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

@Slf4j
@Service
@RequiredArgsConstructor
public class LlmService {

    private final LlmConfig llmConfig;
    private final WebClient.Builder webClientBuilder;

    public Mono<LlmDto.ChatResponse> chat(LlmDto.ChatRequest request) {
        String provider = request.getProvider() != null ? request.getProvider() : "openai";

        if ("custom".equalsIgnoreCase(provider) && isCustomProviderConfigured()) {
            return chatWithCustomProvider(request);
        } else {
            return chatWithOpenAI(request);
        }
    }

    private Mono<LlmDto.ChatResponse> chatWithOpenAI(LlmDto.ChatRequest request) {
        log.info("Sending chat request to OpenAI");

        if (llmConfig.getOpenai().getApiKey() == null || llmConfig.getOpenai().getApiKey().isEmpty()) {
            return Mono.error(new IllegalStateException("OpenAI API密钥未配置"));
        }

        WebClient webClient = webClientBuilder
                .baseUrl(llmConfig.getOpenai().getApiBaseUrl())
                .defaultHeader(HttpHeaders.AUTHORIZATION, "Bearer " + llmConfig.getOpenai().getApiKey())
                .defaultHeader(HttpHeaders.CONTENT_TYPE, MediaType.APPLICATION_JSON_VALUE)
                .build();

        Map<String, Object> requestBody = buildOpenAIRequest(request);

        return webClient.post()
                .uri("/chat/completions")
                .bodyValue(requestBody)
                .retrieve()
                .bodyToMono(Map.class)
                .map(response -> parseOpenAIResponse(response, request))
                .onErrorResume(error -> {
                    log.error("Error calling OpenAI API", error);
                    return Mono.error(new RuntimeException("调用OpenAI API失败: " + error.getMessage()));
                });
    }

    private Mono<LlmDto.ChatResponse> chatWithCustomProvider(LlmDto.ChatRequest request) {
        log.info("Sending chat request to custom LLM provider");

        if (llmConfig.getCustom().getApiUrl() == null || llmConfig.getCustom().getApiUrl().isEmpty()) {
            return Mono.error(new IllegalStateException("自定义LLM端点未配置"));
        }

        WebClient.Builder builder = webClientBuilder
                .baseUrl(llmConfig.getCustom().getApiUrl())
                .defaultHeader(HttpHeaders.CONTENT_TYPE, MediaType.APPLICATION_JSON_VALUE);

        if (llmConfig.getCustom().getApiKey() != null && !llmConfig.getCustom().getApiKey().isEmpty()) {
            builder.defaultHeader(HttpHeaders.AUTHORIZATION, "Bearer " + llmConfig.getCustom().getApiKey());
        }

        WebClient webClient = builder.build();

        // 假设自定义端点使用类似OpenAI的格式
        Map<String, Object> requestBody = buildOpenAIRequest(request);

        return webClient.post()
                .bodyValue(requestBody)
                .retrieve()
                .bodyToMono(Map.class)
                .map(response -> parseCustomResponse(response, request))
                .onErrorResume(error -> {
                    log.error("Error calling custom LLM API", error);
                    return Mono.error(new RuntimeException("调用自定义LLM API失败: " + error.getMessage()));
                });
    }

    private Map<String, Object> buildOpenAIRequest(LlmDto.ChatRequest request) {
        Map<String, Object> requestBody = new HashMap<>();

        // 设置模型
        String model = request.getModel() != null ? request.getModel() : llmConfig.getOpenai().getModel();
        requestBody.put("model", model);

        // 构建消息列表
        List<Map<String, String>> messages = new ArrayList<>();

        // 添加历史消息
        if (request.getHistory() != null) {
            for (LlmDto.Message msg : request.getHistory()) {
                Map<String, String> message = new HashMap<>();
                message.put("role", msg.getRole());
                message.put("content", msg.getContent());
                messages.add(message);
            }
        }

        // 添加当前消息
        Map<String, String> userMessage = new HashMap<>();
        userMessage.put("role", "user");
        userMessage.put("content", request.getMessage());
        messages.add(userMessage);

        requestBody.put("messages", messages);

        // 设置参数
        Integer maxTokens = request.getMaxTokens() != null ? request.getMaxTokens() : llmConfig.getOpenai().getMaxTokens();
        Double temperature = request.getTemperature() != null ? request.getTemperature() : llmConfig.getOpenai().getTemperature();

        requestBody.put("max_tokens", maxTokens);
        requestBody.put("temperature", temperature);

        return requestBody;
    }

    private LlmDto.ChatResponse parseOpenAIResponse(Map<String, Object> response, LlmDto.ChatRequest request) {
        try {
            List<Map<String, Object>> choices = (List<Map<String, Object>>) response.get("choices");
            if (choices == null || choices.isEmpty()) {
                throw new RuntimeException("OpenAI响应中没有choices");
            }

            Map<String, Object> firstChoice = choices.get(0);
            Map<String, String> message = (Map<String, String>) firstChoice.get("message");
            String content = message.get("content");

            Map<String, Object> usage = (Map<String, Object>) response.get("usage");
            Integer totalTokens = usage != null ? (Integer) usage.get("total_tokens") : null;

            String model = (String) response.get("model");

            return LlmDto.ChatResponse.builder()
                    .message(content)
                    .model(model)
                    .tokensUsed(totalTokens)
                    .provider("openai")
                    .build();
        } catch (Exception e) {
            log.error("Error parsing OpenAI response", e);
            throw new RuntimeException("解析OpenAI响应失败: " + e.getMessage());
        }
    }

    private LlmDto.ChatResponse parseCustomResponse(Map<String, Object> response, LlmDto.ChatRequest request) {
        // 这里假设自定义端点使用类似的响应格式
        // 可以根据实际的自定义端点格式进行调整
        return parseOpenAIResponse(response, request);
    }

    private boolean isCustomProviderConfigured() {
        return llmConfig.getCustom().getApiUrl() != null && !llmConfig.getCustom().getApiUrl().isEmpty();
    }
}
